feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path#23523
Open
zhuqi-lucas wants to merge 2 commits into
Open
feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path#23523zhuqi-lucas wants to merge 2 commits into
zhuqi-lucas wants to merge 2 commits into
Conversation
…d schemas on the column-wise path Add `RowsGroupColumn`, a generic `GroupColumn` backed by a single-field `RowConverter` from arrow's row format. Route nested types (`Struct`, `List`, `LargeList`, `FixedSizeList`, and recursive combinations arrow's row converter can encode) to it via `group_column_supported_type` / `make_group_column`. Before this, a single nested column in a GROUP BY key made `supported_schema` return `false` and the whole aggregation fell back to the row-wise `GroupValuesRows`, even when every other column would have qualified for the column-wise fast path. With the generic fallback, only the nested column pays the row-encoding cost; the native columns keep their compact column-wise storage. Impact (measured by `mixed_schema_column_path_uses_less_memory_than_rows_fallback` with 4000 groups of `8 × Int64 + 1 × FixedSizeList<Int64, 4>`): - old (all-rows fallback): 1096 KB - new (column-wise + row-backed nested column): **594 KB (54.2%)** Correctness: - Hashing (`create_hashes` on the raw input columns) already supports nested types, and the caller-side `-0.0/NaN` normalization performed for `GroupValuesRows` keeps hashing aligned with the row-format equality this column implements. Covered by `nested_float_edge_cases_match_rows_fallback`. - Multi-batch streaming intern + `EmitTo::First` / `take_n` covered by `multi_batch_and_emit_first_matches_rows_fallback`. - FSL / Struct roundtrip + `supports_type` invariant covered by the unit tests inside `row_backed`. Follow-ups tracked separately: - Cover types `RowConverter` cannot encode (currently a moving target; if the arrow-rs release DataFusion pins encodes them all — including Map — the fallback is complete). - With full coverage, retire `GroupValuesRows` entirely (motivation: apache#23404). Refs: apache#23404, apache#22715.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves GroupValuesColumn mixed-schema GROUP BY handling by adding a generic row-format-backed GroupColumn for nested types, so only unsupported (nested) columns pay row-encoding overhead instead of forcing the entire key onto the GroupValuesRows fallback.
Changes:
- Introduces
RowsGroupColumn, a genericGroupColumnbacked by Arrow’s row format for nested (row-encodable) types. - Updates the
group_column_supported_type/make_group_columndispatch so nested types useRowsGroupColumnwhile intentionally excluded scalar types remain onGroupValuesRows. - Adds unit and integration-style tests validating correctness (including float edge cases) and reduced memory usage for mixed schemas.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| datafusion/physical-plan/src/aggregates/group_values/row.rs | Documents and exposes dictionary_encode_if_necessary for reuse by the new row-backed column path. |
| datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs | Adds RowsGroupColumn implementation and unit tests for nested types via Arrow row encoding. |
| datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs | Wires nested-type fallback into supported-type checks and factory dispatch; adds mixed-schema memory/correctness tests. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+111
to
+123
| fn rows_to_array<'a>( | ||
| &self, | ||
| rows: impl IntoIterator<Item = arrow::row::Row<'a>>, | ||
| ) -> ArrayRef { | ||
| let mut arrays = self | ||
| .row_converter | ||
| .convert_rows(rows) | ||
| .expect("row conversion during emit"); | ||
| debug_assert_eq!(arrays.len(), 1, "single-field row converter"); | ||
| let array = arrays.swap_remove(0); | ||
| dictionary_encode_if_necessary(&array, &self.output_type) | ||
| .expect("dictionary re-encode during emit") | ||
| } |
Contributor
Author
There was a problem hiding this comment.
Fixed in 132905e — extended the recursion in encode_array_if_necessary to also cover FixedSizeList, LargeList, and Map (parallel to the existing List branch). Added build_preserves_list_of_dictionary_schema as a regression test.
Side note: I discovered while writing tests that RowConverter::convert_rows currently rejects FixedSizeList<Dict> in arrow-rs itself (before our helper runs), so that specific nesting is unreachable through supports_type today. The new FSL branch is defensive coverage for parity with List in case that limitation lifts.
Contributor
Author
There was a problem hiding this comment.
Small correction to my previous reply — I was imprecise about what supports_type does for FixedSizeList<Dict>.
Actual behavior in current arrow-rs: supports_datatype recurses through FixedSizeList into its inner field and hits _ if !d.is_nested() on the Dict leaf, which returns true — so RowsGroupColumn::supports_type(FixedSizeList<Dict<Int32,Utf8>>) does return true. What actually fails is later, inside RowConverter::convert_rows, which panics with "FixedSizeListArray expected data type Dictionary(Int32,Utf8) got Utf8 for 'item'" because arrow-rs constructs the outer FSL against the declared element field but produces flat Utf8 values.
So more accurately: FSL/LargeList/Map pass supports_type but blow up at convert_rows in this arrow-rs version. The FSL / LargeList / Map branches in encode_array_if_necessary are defensive coverage for parity with List — they'll do the right thing once arrow-rs's roundtrip is fixed. The regression test uses List<Dict> because that's the currently-roundtrippable path.
Contributor
|
nice! excited to take a look at this |
Contributor
I think speed is an important piece of this PR — could we at least run the existing benchmarks? IIRC the TPC-H/ClickBench/etc benchmarks don't actually include certain data types, such as dictionaries and FixedSizeList, so they may not be the best indicator of performance. But they should at least give us a ballpark sense of whether this causes a regression. |
Rich-T-kid
suggested changes
Jul 13, 2026
Rich-T-kid
left a comment
Rich-T-kid
left a comment
Contributor
There was a problem hiding this comment.
I took a first pass through, mostly looks good. I left some comments. would be happy to take another look.
| /// Re-apply dictionary / run-end encoding to `array` so it matches `expected`. | ||
| /// | ||
| /// Arrow's [`RowConverter`] decodes dictionary and run-end-encoded values to | ||
| /// their plain value type on the way out, so any group-value array produced |
Contributor
There was a problem hiding this comment.
Arrow's [
RowConverter] decodes dictionary and run-end-encoded values to
their plain value type on the way out
nit: RowConverter decodes dictionarys/REE on their way in since it occurs onRowConverter::append().
Contributor
Author
There was a problem hiding this comment.
You're right, thanks. Fixed the doc in 132905e to say 'flattens dict/REE values during row encoding (at append)' instead of 'on the way out'.
| /// being returned. Shared with the generic row-backed `GroupColumn`. | ||
| /// | ||
| /// [`RowConverter`]: arrow::row::RowConverter | ||
| pub(crate) fn dictionary_encode_if_necessary( |
Contributor
There was a problem hiding this comment.
the comment above mentions casting arrays to REE but the functions is called Dictionary_encode_if_necessary.
Suggested change
| pub(crate) fn dictionary_encode_if_necessary( | |
| pub(crate) fn encode_array_if_necessary( |
or something of the sort
Contributor
Author
There was a problem hiding this comment.
Renamed to encode_array_if_necessary in 132905e.
Comment on lines
+273
to
+275
| /// their plain value type on the way out, so any group-value array produced | ||
| /// from the row format must be re-encoded to the schema's expected type before | ||
| /// being returned. Shared with the generic row-backed `GroupColumn`. |
Contributor
There was a problem hiding this comment.
I think the cast here shows that Dictionaries and REE arrays should have their own specialized implementation of GroupColumn. For now I think its fine but unlike other implementations like fixedSizeList that can go from
arrow-array -> RowFormat -> arrow-array
dictionary/REE need an extra cast step
arrow-array -> rowFormat -> arrow-array -> (cast) -> arrow-array
Contributor
Author
There was a problem hiding this comment.
Agreed — filing as a follow-up so this PR stays scoped to 'keep mixed nested schemas on the column-wise path'. A dedicated DictGroupColumn / ReeGroupColumn would skip the row-encode → cast round-trip entirely, which is a nice next step.
Contributor
There was a problem hiding this comment.
#23187 implements the dictionary version of this if your interested
Comment on lines
+931
to
+936
| // encode them. Gate the fallback to nested types so intentionally-excluded | ||
| // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the | ||
| // `group_column_supported_type` ⇔ `make_group_column` invariant holds. | ||
| if data_type.is_nested() { | ||
| return RowsGroupColumn::supports_type(data_type); | ||
| } |
Contributor
There was a problem hiding this comment.
why are Float16 & Decimal256 arrays not supported for this optimization?
Contributor
Author
There was a problem hiding this comment.
make_group_column doesn't have a type-specialized PrimitiveGroupValueBuilder branch for them — they'd need Float16Type / Decimal256Type wired through the same pattern as Float32Type / Int32Type / etc. We could alternatively route them through RowsGroupColumn by relaxing the is_nested() gate, but I kept this PR scoped to nested types so the group_column_supported_type ⇔ make_group_column biconditional stays clean. Happy to file a follow-up either way if that's useful.
Comment on lines
+1380
to
+1384
| assert!( | ||
| column_size < rows_size, | ||
| "expected column-wise path ({column_size}) to use less memory than \ | ||
| the all-rows fallback ({rows_size})" | ||
| ); |
Contributor
There was a problem hiding this comment.
this is a nice way to guarantee groupColumns will always consume less memory than the GroupValueRows implementation!
Comment on lines
+1396
to
+1414
| /// Relabel a group-index vector so labels are assigned in order of first | ||
| /// appearance. Two vectors are equivalent groupings iff their canonical | ||
| /// forms are equal — this ignores the (opaque, non-semantic) difference in | ||
| /// group-index numbering between the vectorized column path and the | ||
| /// sequential rows fallback. | ||
| fn canonical_grouping(groups: &[usize]) -> Vec<usize> { | ||
| let mut map = HashMap::new(); | ||
| let mut next = 0usize; | ||
| groups | ||
| .iter() | ||
| .map(|&g| { | ||
| *map.entry(g).or_insert_with(|| { | ||
| let v = next; | ||
| next += 1; | ||
| v | ||
| }) | ||
| }) | ||
| .collect() | ||
| } |
Contributor
There was a problem hiding this comment.
This took me a bit of time to understand why its needed. from my understanding its because GroupValues::intern() doesn't specify the order in which new IDs are given. docs
I think the comment is fine as it is but maybe adding something like
GroupValuesimplementations only guarantee that equal rows receive
/// equal group ids and new rows receive a fresh id — the order in which
/// new ids are handed out is not part of the contract, and can differ
/// between correct implementations
can help readers know why this function has to exist. This is mostly a nit
Contributor
Author
There was a problem hiding this comment.
Good point — expanded the doc in 132905e to spell out that GroupValues::intern deliberately does not fix the order of fresh group-ids, and canonicalizing before comparison is what lets us assert equivalence between the vectorized column path and the sequential rows fallback.
Comment on lines
+35
to
+36
| //! and only pays the row-encoding cost on `struct_col`, instead of dragging the | ||
| //! entire key onto `GroupValuesRows`. |
Contributor
There was a problem hiding this comment.
nit:
Suggested change
| //! and only pays the row-encoding cost on `struct_col`, instead of dragging the | |
| //! entire key onto `GroupValuesRows`. | |
| //! and only pays the row-encoding cost on `struct_col`, instead of dragging both | |
| //! columns onto `GroupValuesRows`. |
Comment on lines
+115
to
+118
| let mut arrays = self | ||
| .row_converter | ||
| .convert_rows(rows) | ||
| .expect("row conversion during emit"); |
Contributor
There was a problem hiding this comment.
You may be interested in these issues/ PR in arrow-rs
Contributor
Author
There was a problem hiding this comment.
They're about skipping utf8 validation in RowParser, not about the dict/REE re-encoding path. The connection to this PR is that RowsGroupColumn runs everything through RowConverter (built on RowParser), so once the validate_utf8 = false opt-out lands upstream we could plumb it through for a small row-encode speedup on string / binary group keys. Won't block on it here — good to keep on the radar though, thanks for the pointer.
| } | ||
|
|
||
| fn take_n(&mut self, n: usize) -> ArrayRef { | ||
| debug_assert!(n <= self.group_values.num_rows()); |
Contributor
There was a problem hiding this comment.
does DataFusion guarantee that it wont call GroupValues::emit(emit::to(n)) where n is possibly greater than the number of groups?
Contributor
Author
There was a problem hiding this comment.
Yes. aggregates/order/mod.rs:80 always emits EmitTo::First(n.min(max)), and every other GroupColumn::take_n impl (bytes.rs, bytes_view.rs, boolean.rs, primitive.rs) has the same debug_assert!(self.len() >= n). So the assert is defensive documentation of the caller-side invariant, matching the existing convention.
Contributor
Author
There was a problem hiding this comment.
Correcting my previous reply — I overstated things and re-checked:
- Only
bytes.rsandbytes_view.rshave the explicitdebug_assert!(self.len() >= n);boolean.rsandprimitive.rsdon't (they just rely onVec-slice /BooleanBuffer-slice semantics which panic on out-of-bounds). - The
n.min(max)cap is only in theGroupOrdering::Partial | Fullarm ofoom_emit_to; theGroupOrdering::Nonearm passesnthrough unchecked (order/mod.rs:77).
So the actual guarantee is caller-side: the aggregate stream computes n from either the number of finalized-order groups (Partial/Full) or the group table's own row count (None), both of which are bounded by the values we've interned so far. The debug_assert! in take_n is a defensive documentation of that invariant, consistent with the existing bytes.rs / bytes_view.rs pattern; it's not a universal convention across every GroupColumn impl. Happy to drop it if you'd prefer we match the leaner style in boolean.rs / primitive.rs.
Contributor
There was a problem hiding this comment.
I think its fine as it is
Contributor
Author
Thanks @Rich-T-kid for review, will look into your review comments and address it. |
Fixes from @Rich-T-kid and @copilot-pull-request-reviewer: * row.rs: correct the doc on `encode_array_if_necessary` — arrow's `RowConverter` flattens dictionary / REE values at `append()` (during row encoding), not "on the way out". * row.rs: extend the recursion in `encode_array_if_necessary` to cover `FixedSizeList`, `LargeList`, and `Map` in addition to `Struct` and `List`. Without this, a nested key like `List<Dict<Int32, Utf8>>` emits as `List<Utf8>` because the fallback arm just clones the RowConverter-flattened array, giving a data-type mismatch against the declared output type. * row.rs: rename `dictionary_encode_if_necessary` → `encode_array_if_necessary` since the helper also handles run-end encoding (and now the nested container walks around them). * row_backed.rs: add `build_preserves_list_of_dictionary_schema` — a round-trip regression test that appends a `List<Dict<Int32, Utf8>>` row via `RowsGroupColumn` and asserts the built array's data type matches the declared schema (would fail with the pre-fix helper). * multi_group_by/mod.rs: expand the doc on `canonical_grouping` to spell out that the `GroupValues` trait deliberately does not fix the order in which fresh group-ids are assigned — canonicalizing before comparison is what lets the vectorized column path and the sequential rows fallback be asserted equivalent. * row_backed.rs (module doc): reword "instead of dragging the entire key onto GroupValuesRows" → "instead of dragging both columns onto GroupValuesRows" per Rich-T-kid's suggestion diff.
Contributor
Author
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/rows-backed-group-column (132905e) to 761f7f6 (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/rows-backed-group-column (132905e) to 761f7f6 (merge-base) diff using: tpch File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/rows-backed-group-column (132905e) to 761f7f6 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagedict — base (merge-base)
dict — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
GroupValuesColumnEPIC)GroupValuesRowsentirely (Generic Rows-backed GroupColumn so nested GROUP BY keys stay on the column-wise fast path #23404)Rationale for this change
Today
GroupValuesColumnis all-or-nothing: a single nested column in the GROUP BY key (`Struct`, `List`, `FixedSizeList`, …) makes `supported_schema` return `false` and drops the entire aggregation onto the row-wise `GroupValuesRows` fallback — even when every other column would have qualified for the column-wise fast path. For a `GROUP BY int_col, struct_col` shape, the `int_col` pays the row-encoded storage cost for no reason.What changes are included in this PR?
Add `RowsGroupColumn`: a generic `GroupColumn` backed by a single-field `RowConverter`, wired in as the nested-type dispatch arm of `group_column_supported_type` / `make_group_column`. Native columns keep their type-specialized builders; the nested column pays row-encoding only for its one column.
Gated to `data_type.is_nested()` so intentionally excluded scalar types (Float16, Decimal256) stay on `GroupValuesRows` and the `group_column_supported_type` ⇔ `make_group_column` invariant holds.
Impact
Memory, measured with 4000 groups of `8 × Int64 + 1 × FixedSizeList<Int64, 4>` in `mixed_schema_column_path_uses_less_memory_than_rows_fallback`:
Speed: not benchmarked as a headline result — the wins come from native columns keeping their type-specialized `equal_to`/`append_val` fast paths instead of falling back to byte-encoded row comparisons.
Are these changes tested?
Yes:
All 39 tests in `aggregates::group_values` pass.
Are there any user-facing changes?
No — internal aggregation representation only. Same query results, lower memory footprint on mixed-schema GROUP BY keys.
Follow-ups (out of scope)